import yaml
from .base_command import BaseCommand
from cleo.helpers import option
from tjob.initialize.setup_envs import (
    dump_pai_key_passwd,
    set_deekeeper,
    set_dlc_login,
    _CACHE,
    set_dlc_params,
    set_local_params,
    MACHINE_LIST,
    _LOCAL_CONF,
    # clear_envs,
)
from tjob.initialize.setup_envs_cloud import (
    _CLOUD_ACCESS, _CLOUD_DEFAULT_CONF, _CLOUD_JOB_PARAMS, get_cloud_default_conf, get_cloud_job_default_params, set_cloud_login, set_cloud_default_params
)

from tjob.utils import str2dict
import os.path as osp


class InitCommand(BaseCommand):
    name = "init"
    description = "Initilaize tjob."
    options = [
        option(
            "username",
            "u",
            flag=False,
            description="your pai username",
            value_required=False,
        ),
        option(
            "password",
            "p",
            flag=False,
            description="your pai password",
            value_required=False,
        ),
        option(
            "keeper_access",
            "a",
            flag=False,
            description="your deekpeeper access key",
        ),
        option(
            "keeper_secret",
            "s",
            flag=False,
            description="your deekpeeper secret key",
        ),
        option(
            "default_params",
            "d",
            description="custom your defalut params. format: key1=value1 key2=value2"
            + "e.g. -d data_sources=xahs12 worker_gpu=8",
            flag=False,
        ),
        option(
            "platform",
            description="select default platform. [pai/local/cloud]",
            flag=False,
            value_required=False,
            default="",
            
        ),
    ]

    def _handle_input(self, name: str, msg: str):
        if not self.option(name):
            val = input(msg)
        else:
            val = self.option(name)
        val = val.strip()
        return val
    
    def _handle_cloud_init(self):
        user = self.option("username")
        password = self.option("password")
        if not (user and password):
            self.red_text("[FAILED] user and password cannot be empty. Exit tjob.")
            exit(1)
        set_cloud_login(user, password)
        if self.option("default_params"):
            params = str2dict(self.option("default_params"), sep_ele=" ", sep_value="=")
        else:
            params = {}
        set_cloud_default_params(params)
        
        default_params = get_cloud_job_default_params()
        conf = get_cloud_default_conf()
        self.blue_text("[Init] Done")
        self.raw_text(f"cloud default auth file path: {_CLOUD_ACCESS}")
        with open(_CLOUD_ACCESS, "r") as f:
            self.green_text(f.read())
        self.raw_text(f"cloud default config file path: {_CLOUD_DEFAULT_CONF}")
        with open(_CLOUD_DEFAULT_CONF, "r") as f:
            self.green_text(f.read())
        self.raw_text(f"cloud default job file path: {_CLOUD_JOB_PARAMS}")
        with open(_CLOUD_JOB_PARAMS, "r") as f:
            self.green_text(f.read())

        self.add_style("bb", fg="blue", options=["bold"])
        self.add_style("bg", fg="green", options=["bold"])
        self.add_style("w", fg="white")

        extra_deekeeper_msg = (
            "<bb>Deepkeeper <w> an excellent tool for managing experiments. We recommend using it to <bb>RECORD, DISPLAY, SHARE, BACK-UP<bb> <w> your experiment!\n"  # noqa
            + "<w>Get your Deekeeper access key from: <bg> https://deekeeper-fe.srv.deeproute.cn/ \n"  # noqa
            + "<w>More about Deekeeper please refer to: <bg>https://rqk9rsooi4.feishu.cn/wiki/wikcnMFarrl3BwYE36No6OnSqZd \n"  # noqa
        )

        if not self.option("keeper_access"):
            self.line(extra_deekeeper_msg)
        else:
            access = self.option("keeper_access")
            secret = self.option("keeper_secret")
            if access and secret:
                set_deekeeper(access, secret, self)
        

    def handle(self):
        
        platform = self.option("platform")
        if platform == "cloud":
            self._handle_cloud_init()
            return  
        
        user = self._handle_input("username", "Enter your PAI username:")
        password = self._handle_input("password", "Enter your PAI password:")
        dump_pai_key_passwd(user, password)
        if user != "":
            returncode = set_dlc_login(user, password)
            if returncode != 0:
                self.red_text("[FAILED] init PAI. Check user and passwd. Exit tjob.")
                exit(0)

        
        # note: init args is same as dlc.params not, submit cli
        # this may confuse
        if self.option("default_params"):
            params = str2dict(self.option("default_params"), sep_ele=" ", sep_value="=")
        else:
            params = {}

        set_dlc_params(params)
        set_local_params()

        self.add_style("bb", fg="blue", options=["bold"])
        self.add_style("bg", fg="green", options=["bold"])
        self.add_style("w", fg="white")

        extra_deekeeper_msg = (
            "<bb>Deepkeeper <w> an excellent tool for managing experiments. We recommend using it to <bb>RECORD, DISPLAY, SHARE, BACK-UP<bb> <w> your experiment!\n"  # noqa
            + "<w>Get your Deekeeper access key from: <bg> https://deekeeper-fe.srv.deeproute.cn/ \n"  # noqa
            + "<w>More about Deekeeper please refer to: <bg>https://rqk9rsooi4.feishu.cn/wiki/wikcnMFarrl3BwYE36No6OnSqZd \n"  # noqa
        )

        if not self.option("keeper_access"):
            self.line(extra_deekeeper_msg)

        access = self._handle_input(
            "keeper_access", "Enter your Deekeeper Access Key (q to exit):"
        )
        if access not in ["q", ""]:
            secret = self._handle_input(
                "keeper_secret", "Enter your Deekeeper secret Key:"
            )
            set_deekeeper(access, secret, self)

        self.raw_text("`tjob` init default submit config of PAI in: ")
        for machine in MACHINE_LIST:
            _machine_cfg = _CACHE + f"/{machine}.submit.dlc.yaml"
            if not osp.exists(_machine_cfg):
                raise IOError(f"Initialize Failed at {machine}.")
            self.raw_text("    " + _machine_cfg)
        self.raw_text("`tjob` init default submit config of Local in: ")
        self.raw_text(f"    {_LOCAL_CONF}")

        self.blue_text("[Init] Done")
